import os
import glob
import numpy as np
import pandas as pd
from scipy.stats import skew, kurtosis
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import classification_report

# 读取HRV数据并提取特征

def extract_features_from_df(df, pid=None):
    df.columns = df.columns.str.strip()
    required_cols = ['mean_nni', 'sdnn', 'rmssd', 'tinn']
    missing = [col for col in required_cols if col not in df.columns]
    if missing:
        raise KeyError(f"Patient {pid}: 缺少列 {missing}")
    feats = {}
    for col in required_cols:
        arr = pd.to_numeric(df[col], errors='coerce').dropna().values
        if arr.size == 0:
            raise ValueError(f"Patient {pid}: 列 '{col}' 全部空值")
        feats[f'{col}_mean'] = arr.mean()
        feats[f'{col}_std'] = arr.std(ddof=1)
        feats[f'{col}_range'] = arr.max() - arr.min()
        feats[f'{col}_iqr'] = np.percentile(arr, 75) - np.percentile(arr, 25)
        feats[f'{col}_skew'] = skew(arr) if arr.size > 2 else np.nan
        feats[f'{col}_kurtosis'] = kurtosis(arr) if arr.size > 3 else np.nan
    return feats


def load_psg_info(xlsx_path):
    info = pd.read_excel(xlsx_path)
    info = info.rename(columns={'顺序': 'patient_id'})
    info['patient_id'] = info['patient_id'].astype(str)
    bins = list(range(0, 101, 10))
    labels = [f"{i}-{i+9}" for i in bins[:-1]]
    info['age_bin'] = pd.cut(info['年龄'], bins=bins, labels=labels, right=False)
    return info[['patient_id', '性别', 'age_bin']]


def build_feature_dataframe(hrv_folder):
    all_files = sorted(glob.glob(os.path.join(hrv_folder, '*.csv')))
    records = []
    for fp in all_files:
        pid = os.path.splitext(os.path.basename(fp))[0]
        try:
            df = pd.read_csv(fp, sep=None, engine='python', header=0)
            feats = extract_features_from_df(df, pid)
            feats['patient_id'] = pid
            records.append(feats)
        except Exception as e:
            print(f"⚠️ WARNING: 跳过 Patient {pid}, 原因: {e}")
    if not records:
        raise RuntimeError("未提取到任何特征，请检查 CSV 文件格式和路径！")
    features = pd.DataFrame(records)
    print(f"✅ 成功提取 {len(features)} 位患者特征。")
    return features


def train_and_predict_all(features, info):
    data = features.merge(info, on='patient_id', how='inner')
    X = data.drop(columns=['patient_id', '性别', 'age_bin'])
    # 性别编码与训练
    le_sex = LabelEncoder()
    y_sex = le_sex.fit_transform(data['性别'])
    X_tr_s, X_te_s, y_tr_s, y_te_s = train_test_split(X, y_sex, test_size=0.3, random_state=42)
    clf_sex = RandomForestClassifier(n_estimators=100, random_state=42)
    clf_sex.fit(X_tr_s, y_tr_s)

    # 年龄段编码与训练
    le_age = LabelEncoder()
    y_age = le_age.fit_transform(data['age_bin'].astype(str))
    X_tr_a, X_te_a, y_tr_a, y_te_a = train_test_split(X, y_age, test_size=0.3, random_state=42)
    clf_age = RandomForestClassifier(n_estimators=100, random_state=42)
    clf_age.fit(X_tr_a, y_tr_a)

    # 全样本预测结果
    data['pred_gender'] = le_sex.inverse_transform(clf_sex.predict(X))
    data['pred_age_bin'] = le_age.inverse_transform(clf_age.predict(X))
    data.to_csv('patient_predictions.csv', index=False)
    print("✅ 已保存所有患者预测结果到 patient_predictions.csv")

    # 性别分类报告
    print("=== 性别分类（测试集） ===")
    print(classification_report(
        y_te_s,
        clf_sex.predict(X_te_s),
        labels=list(range(len(le_sex.classes_))),
        target_names=le_sex.classes_,
        zero_division=0
    ))

    # 年龄段分类报告
    print("=== 年龄段分类（测试集） ===")
    print(classification_report(
        y_te_a,
        clf_age.predict(X_te_a),
        labels=list(range(len(le_age.classes_))),
        target_names=le_age.classes_,
        zero_division=0
    ))

    # 交叉验证准确率
    print(f"性别 5 折 CV 准确率：{cross_val_score(clf_sex, X, y_sex, cv=5).mean():.3f}")
    print(f"年龄段 5 折 CV 准确率：{cross_val_score(clf_age, X, y_age, cv=5).mean():.3f}")

if __name__ == '__main__':
    xlsx_path = r"C:\py\sleepbjut-master\data\psg_info.xlsx"
    hrv_folder = r"C:\py\sleepbjut-master\csv\hrv_only"
    info = load_psg_info(xlsx_path)
    features = build_feature_dataframe(hrv_folder)
    train_and_predict_all(features, info)
